package stablehlo

import (
	"github.com/gomlx/gomlx/backends"
	"github.com/gomlx/gopjrt/dtypes"
)

// Capabilities of the SimpleGo backends: the set of supported operations and data types.
var Capabilities = backends.Capabilities{
	Operations: map[backends.OpType]bool{
		// Graph inputs (leaf nodes)
		backends.OpTypeParameter: true,
		backends.OpTypeConstant:  true,

		// Standard unary operations:
		backends.OpTypeAbs:        true,
		backends.OpTypeBitCount:   true,
		backends.OpTypeBitwiseNot: true,
		backends.OpTypeCeil:       true,
		backends.OpTypeClz:        true,
		backends.OpTypeCos:        true,
		backends.OpTypeErf:        true,
		backends.OpTypeExp:        true,
		backends.OpTypeExpm1:      true,
		backends.OpTypeFloor:      true,
		backends.OpTypeIsFinite:   true,
		backends.OpTypeIsNaN:      true,
		backends.OpTypeLog1p:      true,
		backends.OpTypeLog:        true,
		backends.OpTypeLogicalNot: true,
		backends.OpTypeLogistic:   true,
		backends.OpTypeNeg:        true,
		backends.OpTypeRound:      true,
		backends.OpTypeRsqrt:      true,
		backends.OpTypeSign:       true,
		backends.OpTypeSin:        true,
		backends.OpTypeSqrt:       true,
		backends.OpTypeTanh:       true,

		// Standard binary operations:
		backends.OpTypeAdd:        true,
		backends.OpTypeBitwiseAnd: true,
		backends.OpTypeBitwiseOr:  true,
		backends.OpTypeBitwiseXor: true,
		backends.OpTypeDiv:        true,
		backends.OpTypeLogicalAnd: true,
		backends.OpTypeLogicalOr:  true,
		backends.OpTypeLogicalXor: true,
		backends.OpTypeMax:        true,
		backends.OpTypeMin:        true,
		backends.OpTypeMul:        true,
		backends.OpTypePow:        true,
		backends.OpTypeRem:        true,
		backends.OpTypeSub:        true,

		// Comparison operators.
		backends.OpTypeEqual:                    true,
		backends.OpTypeEqualTotalOrder:          true,
		backends.OpTypeGreaterOrEqual:           true,
		backends.OpTypeGreaterOrEqualTotalOrder: true,
		backends.OpTypeGreaterThan:              true,
		backends.OpTypeGreaterThanTotalOrder:    true,
		backends.OpTypeLessOrEqual:              true,
		backends.OpTypeLessOrEqualTotalOrder:    true,
		backends.OpTypeLessThan:                 true,
		backends.OpTypeLessThanTotalOrder:       true,
		backends.OpTypeNotEqual:                 true,
		backends.OpTypeNotEqualTotalOrder:       true,

		// Complex operations:
		backends.OpTypeComplex: true,
		backends.OpTypeConj:    true,
		backends.OpTypeImag:    true,
		backends.OpTypeReal:    true,

		// Other operations:
		backends.OpTypeArgMinMax:             true,
		backends.OpTypeBatchNormForInference: true,
		backends.OpTypeBatchNormForTraining:  true,
		backends.OpTypeBatchNormGradient:     true,
		backends.OpTypeBitcast:               true,
		backends.OpTypeBroadcast:             true,
		backends.OpTypeBroadcastInDim:        true,
		backends.OpTypeClamp:                 true,
		backends.OpTypeConcatenate:           true,
		backends.OpTypeConvertDType:          true,
		backends.OpTypeConvGeneral:           true,
		backends.OpTypeDynamicSlice:          true,
		backends.OpTypeDynamicUpdateSlice:    true,
		backends.OpTypeDot:                   true,
		backends.OpTypeDotGeneral:            true,
		backends.OpTypeFFT:                   true,
		backends.OpTypeGather:                true,
		backends.OpTypeIdentity:              true,
		backends.OpTypeIota:                  true,
		backends.OpTypePad:                   true,
		backends.OpTypeReduceBitwiseAnd:      true,
		backends.OpTypeReduceBitwiseOr:       true,
		backends.OpTypeReduceBitwiseXor:      true,
		backends.OpTypeReduceLogicalAnd:      true,
		backends.OpTypeReduceLogicalOr:       true,
		backends.OpTypeReduceLogicalXor:      true,
		backends.OpTypeReduceMax:             true,
		backends.OpTypeReduceMin:             true,
		backends.OpTypeReduceProduct:         true,
		backends.OpTypeReduceSum:             true,
		backends.OpTypeReduceWindow:          true,
		backends.OpTypeReshape:               true,
		backends.OpTypeReverse:               true,
		backends.OpTypeRNGBitGenerator:       true,
		backends.OpTypeScatterSum:            true,
		backends.OpTypeScatterMax:            true,
		backends.OpTypeScatterMin:            true,
		backends.OpTypeSelectAndScatterMax:   true,
		backends.OpTypeSelectAndScatterMin:   true,
		backends.OpTypeShiftLeft:             true,
		backends.OpTypeShiftRightArithmetic:  true,
		backends.OpTypeShiftRightLogical:     true,
		backends.OpTypeSlice:                 true,
		backends.OpTypeTranspose:             true,
		backends.OpTypeWhere:                 true,

		// Collective (distributed across devices) operations:
		backends.OpTypeAllReduce: true,
	},

	DTypes: map[dtypes.DType]bool{
		dtypes.Bool:       true,
		dtypes.Int8:       true,
		dtypes.Int16:      true,
		dtypes.Int32:      true,
		dtypes.Int64:      true,
		dtypes.Uint8:      true,
		dtypes.Uint16:     true,
		dtypes.Uint32:     true,
		dtypes.Uint64:     true,
		dtypes.Float32:    true,
		dtypes.Float64:    true,
		dtypes.BFloat16:   true,
		dtypes.Complex64:  true,
		dtypes.Complex128: true,
	},
}
